#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import sys
import os
import pickle
import numpy as np
import re # Added for improved JSON parsing
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config.config import settings

try:
    from sentence_transformers import SentenceTransformer
    import faiss
    SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError:
    SENTENCE_TRANSFORMERS_AVAILABLE = False
    print("⚠️ 请先安装依赖: pip install sentence-transformers faiss-cpu")

# 尝试导入智谱AI SDK
try:
    from zhipuai import ZhipuAI
    ZHIPU_AVAILABLE = True
except ImportError:
    ZHIPU_AVAILABLE = False
    print("⚠️ 智谱AI SDK未安装，请运行: pip install zhipuai")

class RAGInvalidValueGenerator:
    """RAG检索增强的异常值生成器"""
    
    def __init__(self):
        self.chunks = None
        self.index = None
        self.model = None
        self.load_rag_index()
    
    def load_rag_index(self):
        """加载RAG索引"""
        if not SENTENCE_TRANSFORMERS_AVAILABLE:
            print("❌ 缺少sentence-transformers依赖")
            return False
        
        if not os.path.exists('data/rag_chunks.pkl') or not os.path.exists('data/rag_index.faiss'):
            print("❌ RAG索引文件不存在，请先运行 scripts/build_rag_index.py")
            return False
        
        try:
            # 加载分段
            with open('data/rag_chunks.pkl', 'rb') as f:
                self.chunks = pickle.load(f)
            
            # 加载索引
            self.index = faiss.read_index('data/rag_index.faiss')
            
            # 尝试使用本地模型路径
            model_path = "./models/sentence-transformers"
            if os.path.exists(model_path):
                print(f"🔍 使用本地模型: {model_path}")
                self.model = SentenceTransformer(model_path)
            else:
                print("🔍 使用在线模型: paraphrase-multilingual-MiniLM-L12-v2")
                self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
            
            print(f"✅ RAG索引加载成功，包含 {len(self.chunks)} 个文档片段")
            return True
            
        except Exception as e:
            print(f"❌ 加载RAG索引失败: {e}")
            return False
    
    def retrieve(self, query, top_k=3):
        """检索相关文档片段"""
        if not self.chunks or not self.index or not self.model:
            print("❌ RAG索引未正确加载")
            return []
        
        try:
            # 向量化查询
            q_emb = self.model.encode([query])
            
            # 检索
            D, I = self.index.search(np.array(q_emb).astype('float32'), top_k)
            
            # 返回相关片段
            relevant_chunks = []
            for i, (dist, idx) in enumerate(zip(D[0], I[0])):
                chunk = self.chunks[idx]
                similarity = 1 - dist
                relevant_chunks.append({
                    'content': chunk,
                    'similarity': similarity,
                    'rank': i + 1
                })
            
            return relevant_chunks
            
        except Exception as e:
            print(f"❌ 检索失败: {e}")
            return []
    
    def call_ai_model(self, prompt):
        """调用AI模型"""
        try:
            ai_model_type = settings.AI_MODEL_TYPE.lower()
            
            if ai_model_type == "zhipu":
                return self.call_zhipu_model(prompt)
            elif ai_model_type == "openai":
                return self.call_openai_model(prompt)
            elif ai_model_type == "qwen":
                return self.call_qwen_model(prompt)
            elif ai_model_type == "doubao":
                return self.call_doubao_model(prompt)
            else:
                print(f"⚠️ 不支持的AI模型类型: {ai_model_type}")
                return None
                
        except Exception as e:
            print(f"⚠️ AI模型调用异常: {e}")
            return None
    
    def call_zhipu_model(self, prompt):
        """调用智谱AI模型"""
        try:
            api_key = settings.ZHIPU_API_KEY
            model = settings.ZHIPU_MODEL
            
            if not api_key:
                print("⚠️ 未配置ZHIPU_API_KEY")
                return None
            
            if not ZHIPU_AVAILABLE:
                print("⚠️ 智谱AI SDK未安装")
                return None
            
            client = ZhipuAI(api_key=api_key)
            
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                max_tokens=2000,
                temperature=0.7
            )
            
            return response.choices[0].message.content.strip()
                
        except Exception as e:
            print(f"⚠️ 智谱AI模型调用异常: {e}")
            return None
    
    def call_openai_model(self, prompt):
        """调用OpenAI模型"""
        try:
            api_key = settings.OPENAI_API_KEY
            model = settings.OPENAI_MODEL
            
            if not api_key:
                print("⚠️ 未配置OPENAI_API_KEY")
                return None
            
            import requests
            
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": model,
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个API测试数据生成专家，专门生成各种字段类型的异常值用于测试。"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "max_tokens": 2000,
                "temperature": 0.7
            }
            
            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result["choices"][0]["message"]["content"].strip()
            else:
                print(f"⚠️ OpenAI API调用失败: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"⚠️ OpenAI模型调用异常: {e}")
            return None
    
    def call_qwen_model(self, prompt):
        """调用通义千问模型"""
        try:
            api_key = settings.QWEN_API_KEY
            model = settings.QWEN_MODEL
            
            if not api_key:
                print("⚠️ 未配置QWEN_API_KEY")
                return None
            
            import requests
            
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": model,
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个API测试数据生成专家，专门生成各种字段类型的异常值用于测试。"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "max_tokens": 2000,
                "temperature": 0.7
            }
            
            response = requests.post(
                "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result["output"]["text"].strip()
            else:
                print(f"⚠️ 通义千问API调用失败: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"⚠️ 通义千问模型调用异常: {e}")
            return None
    
    def call_doubao_model(self, prompt):
        """调用豆包AI模型"""
        try:
            api_key = settings.DOUBAO_API_KEY
            model = settings.DOUBAO_MODEL
            base_url = settings.DOUBAO_BASE_URL
            
            if not api_key:
                print("⚠️ 未配置DOUBAO_API_KEY")
                return None
            
            import requests
            
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": model,
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个API测试数据生成专家，专门生成各种字段类型的异常值用于测试。"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "max_tokens": 2000,
                "temperature": 0.7
            }
            
            response = requests.post(
                base_url,
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result["choices"][0]["message"]["content"].strip()
            else:
                print(f"⚠️ 豆包AI API调用失败: {response.status_code}")
                return None
                
        except Exception as e:
            print(f"⚠️ 豆包AI模型调用异常: {e}")
            return None
    
    def parse_ai_response(self, ai_response):
        """解析AI返回的内容，提取JSON数组"""
        if not ai_response:
            return []
        
        try:
            # 清理响应内容
            cleaned_response = ai_response.strip()
            
            # 移除markdown代码块标记
            cleaned_response = re.sub(r'```json\s*', '', cleaned_response)
            cleaned_response = re.sub(r'```\s*', '', cleaned_response)
            
            # 移除可能的解释文字
            lines = cleaned_response.split('\n')
            json_lines = []
            in_json = False
            
            for line in lines:
                line = line.strip()
                if line.startswith('[') or line.startswith('"'):
                    in_json = True
                if in_json:
                    json_lines.append(line)
                if line.endswith(']'):
                    break
            
            if json_lines:
                cleaned_response = '\n'.join(json_lines)
            
            # 尝试直接解析
            try:
                invalid_values = json.loads(cleaned_response)
                if isinstance(invalid_values, list):
                    return invalid_values
            except json.JSONDecodeError:
                pass
            
            # 如果直接解析失败，尝试提取数组
            array_match = re.search(r'\[.*?\]', cleaned_response, re.DOTALL)
            if array_match:
                try:
                    array_str = array_match.group(0)
                    invalid_values = json.loads(array_str)
                    if isinstance(invalid_values, list):
                        return invalid_values
                except json.JSONDecodeError:
                    pass
            
            # 如果还是失败，尝试手动解析
            return self.parse_array_manually(cleaned_response)
            
        except Exception as e:
            print(f"❌ JSON解析失败: {e}")
            return []
    
    def parse_array_manually(self, text):
        """手动解析数组，处理格式不标准的JSON"""
        try:
            # 查找数组开始和结束
            start = text.find('[')
            end = text.rfind(']')
            
            if start == -1 or end == -1 or start >= end:
                return []
            
            array_content = text[start+1:end]
            
            # 分割字符串，处理引号内的逗号
            values = []
            current_value = ""
            in_quotes = False
            escape_next = False
            
            for char in array_content:
                if escape_next:
                    current_value += char
                    escape_next = False
                elif char == '\\':
                    escape_next = True
                    current_value += char
                elif char == '"' and not escape_next:
                    in_quotes = not in_quotes
                    current_value += char
                elif char == ',' and not in_quotes:
                    # 分割点
                    value = current_value.strip().strip('"')
                    if value:
                        values.append(value)
                    current_value = ""
                else:
                    current_value += char
            
            # 添加最后一个值
            if current_value.strip():
                value = current_value.strip().strip('"')
                if value:
                    values.append(value)
            
            return values
            
        except Exception as e:
            print(f"❌ 手动解析失败: {e}")
            return []

    def generate_invalid_values_for_field(self, field_name, field_type):
        """为特定字段生成异常值"""
        print(f"\n🔍 正在为字段 {field_name} ({field_type}) 生成异常值...")
        
        # 构建检索查询
        query = f"{field_name}字段的格式限制和异常值规则"
        
        # 检索相关文档片段
        relevant_chunks = self.retrieve(query, top_k=2)
        
        if not relevant_chunks:
            print(f"⚠️ 未找到 {field_name} 字段的相关信息")
            return []
        
        print(f"📄 检索到 {len(relevant_chunks)} 个相关文档片段:")
        for chunk in relevant_chunks:
            print(f"  相似度: {chunk['similarity']:.3f}")
            print(f"  内容: {chunk['content'][:100]}...")
        
        # 构建prompt
        context = "\n\n".join([chunk['content'] for chunk in relevant_chunks])
        
        prompt = f"""
你是一个API测试数据生成专家。以下是关于{field_name}字段的规范文档内容：

{context}

请根据以上规范文档，为{field_name}字段生成10个异常值，覆盖以下错误类型：
1. 格式错误（如格式不匹配、结构错误等）
2. 长度超限（如超出最大长度、不足最小长度等）
3. 类型错误（如类型不匹配、值类型错误等）
4. 缺失必填字段（如缺少必要字段等）
5. 无效值（如超出范围、无效选项等）
6. 边界值错误（如边界值、特殊字符等）

请严格按照以下JSON数组格式返回，不要添加任何解释文字：

["异常值1", "异常值2", "异常值3", "异常值4", "异常值5", "异常值6", "异常值7", "异常值8", "异常值9", "异常值10"]

重要提示：
- 只返回JSON数组，不要包含任何其他文字
- 确保每个异常值都用双引号包围
- 确保JSON格式完全正确
- 对于复杂类型，异常值应该是完整的JSON字符串
"""

        # 调用AI生成
        print(f"🤖 正在调用AI生成异常值...")
        ai_response = self.call_ai_model(prompt)
        
        if ai_response:
            print(f"✅ AI返回内容:")
            print(f"---AI开始---")
            print(ai_response)
            print(f"---AI结束---")
            
            # 使用改进的解析方法
            invalid_values = self.parse_ai_response(ai_response)
            
            if invalid_values:
                print(f"✅ 成功生成 {len(invalid_values)} 个异常值")
                return invalid_values
            else:
                print(f"❌ 无法解析AI返回的内容")
                return []
        else:
            print(f"❌ AI调用失败")
            return []
    
    def generate_all_invalid_values(self):
        """生成所有字段的异常值"""
        print("=== RAG检索增强的异常值生成器 ===")
        
        # 定义字段列表 - 包含文档中定义的所有12个字段
        fields = [
            ("text_7ea7d6a3339", "text"),
            ("bigint_7d6a33398bf", "bigint"),
            ("number_a33398bfae5", "number"),
            ("date_398bfae52d8", "date"),
            ("datetime_bfae52d8694", "datetime"),
            ("option_e52d86948bb", "option"),
            ("boolean_e206a267520", "boolean"),
            ("richText_76d22d35024", "richText"),
            ("multilingual_24673c280f2", "multilingual"),
            ("phone_73c280f2783", "mobileNumber"),  # 字段名是phone，但类型是mobileNumber
            ("email_f2783b49d9a", "email"),
            ("decimal_a446e77c109", "decimal")
        ]
        
        print(f"📋 将处理 {len(fields)} 个字段:")
        for i, (field_name, field_type) in enumerate(fields, 1):
            print(f"  {i:2d}. {field_name} ({field_type})")
        
        all_invalid_values = {}
        
        for field_name, field_type in fields:
            invalid_values = self.generate_invalid_values_for_field(field_name, field_type)
            if invalid_values:
                # 使用具体的接口名作为键，而不是字段类型
                all_invalid_values[field_name] = invalid_values
        
        return all_invalid_values

def main():
    """主函数"""
    print("=" * 60)
    print("RAG检索增强的异常值生成器")
    print("=" * 60)
    
    # 创建生成器
    generator = RAGInvalidValueGenerator()
    
    # 生成异常值
    invalid_values = generator.generate_all_invalid_values()
    
    if invalid_values:
        # 保存到文件
        output_file = "data/rag_invalidCollection.json"
        
        # 确保目录存在
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(invalid_values, f, 
                     ensure_ascii=False,
                     indent=2,
                     sort_keys=True)
        
        print(f"\n✅ 成功生成异常值集合并保存到 {output_file}")
        
        # 统计信息
        print(f"\n📊 异常值集合统计:")
        for field_name, values in invalid_values.items():
            print(f"  {field_name}: {len(values)} 个异常值")
        
        # 显示生成的JSON结构预览
        print(f"\n📋 生成的JSON结构预览:")
        print("```json")
        preview = {}
        for field_name, values in list(invalid_values.items())[:3]:  # 只显示前3个
            preview[field_name] = values[:3]  # 每个字段只显示前3个值
        print(json.dumps(preview, ensure_ascii=False, indent=2))
        print("```")
        
    else:
        print("❌ 未能生成任何异常值")

if __name__ == "__main__":
    main() 